package com.intelligent.system.common.sse;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class SseEmitterManager {
    private static final Map<String, SseEmitter> SSE_CACHE = new ConcurrentHashMap<>();
    private static final long DEFAULT_TIMEOUT = 60 * 1000L;
    private static final long RECONNECT_TIMEOUT = 30 * 1000L;
    private final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 创建SSE连接
     */
    public SseEmitter createConnection(String clientId) {
        // 设置超时时间为1分钟
        SseEmitter sseEmitter = new SseEmitter(DEFAULT_TIMEOUT);
        
        // 注册回调
        sseEmitter.onCompletion(() -> {
            log.info("客户端 {} 的SSE连接已完成", clientId);
            SSE_CACHE.remove(clientId);
        });
        
        sseEmitter.onTimeout(() -> {
            log.info("客户端 {} 的SSE连接超时", clientId);
            SSE_CACHE.remove(clientId);
        });
        
        sseEmitter.onError(throwable -> {
            log.error("客户端 {} 的SSE连接发生错误", clientId, throwable);
            SSE_CACHE.remove(clientId);
        });

        // 缓存连接
        SSE_CACHE.put(clientId, sseEmitter);
        log.info("已为客户端 {} 创建SSE连接", clientId);
        
        // 发送连接成功消息
        try {
            SseMessage connectMessage = new SseMessage("connect", "Connected successfully", clientId);
            sseEmitter.send(objectMapper.writeValueAsString(connectMessage));
        } catch (IOException e) {
            log.error("向客户端 {} 发送连接消息失败", clientId, e);
        }
        
        return sseEmitter;
    }

    /**
     * 发送消息给指定客户端
     */
    public void sendMessage(String clientId, SseMessage message) {
        SseEmitter sseEmitter = SSE_CACHE.get(clientId);
        if (sseEmitter != null) {
            try {
                sseEmitter.send(objectMapper.writeValueAsString(message));
            } catch (IOException e) {
                log.error("向客户端 {} 发送消息失败", clientId, e);
                SSE_CACHE.remove(clientId);
            }
        }
    }

    /**
     * 发送消息给所有客户端
     */
    public void sendMessageToAll(SseMessage message) {
        SSE_CACHE.forEach((clientId, emitter) -> {
            try {
                message.setClientId(clientId);
                emitter.send(objectMapper.writeValueAsString(message));
            } catch (IOException e) {
                log.error("向客户端 {} 发送消息失败", clientId, e);
                SSE_CACHE.remove(clientId);
            }
        });
    }

    /**
     * 关闭指定客户端的连接
     */
    public void closeConnection(String clientId) {
        SseEmitter sseEmitter = SSE_CACHE.get(clientId);
        if (sseEmitter != null) {
            sseEmitter.complete();
            SSE_CACHE.remove(clientId);
            log.info("已关闭客户端 {} 的SSE连接", clientId);
        }
    }

    /**
     * 关闭所有连接
     */
    public void closeAllConnections() {
        SSE_CACHE.forEach((clientId, emitter) -> {
            emitter.complete();
            log.info("已关闭用户 {} 的SSE连接", clientId);
        });
        SSE_CACHE.clear();
    }

    /**
     * 获取在线客户端数量
     */
    public int getOnlineCount() {
        return SSE_CACHE.size();
    }
} 